import osmnx as ox
import geopandas as gpd


def get_road_network(place_name, network_type='all'):
    """
    获取指定地区的道路网络

    参数:
        place_name (str): 地区名称
        network_type (str): 道路网络类型，可选值包括 'all', 'drive', 'walk', 'bike'

    返回:
        nx.MultiDiGraph: 道路网络图
    """
    # 设置osmnx配置
    ox.config(use_cache=True, log_console=True)

    try:
        # 获取道路网络
        G = ox.graph_from_place(place_name, network_type=network_type)
        return G
    except Exception as e:
        print(f"获取道路网络时出错: {str(e)}")
        return None


def graph_to_geojson(G, output_path):
    """
    将道路网络转换为GeoJSON格式并保存

    参数:
        G (nx.MultiDiGraph): 道路网络图
        output_path (str): 输出文件路径
    """
    try:
        # 将网络转换为GeoDataFrame
        nodes, edges = ox.graph_to_gdfs(G)

        # 保存为GeoJSON
        edges.to_file(output_path, driver='GeoJSON')
        print(f"成功保存道路网络到: {output_path}")
    except Exception as e:
        print(f"保存GeoJSON时出错: {str(e)}")


def main():
    # 定义地区名称
    place_name = "Yuzhong, Chongqing, China"

    # 获取道路网络
    print(f"正在获取 {place_name} 的道路网络...")
    G = get_road_network(place_name)

    if G is not None:
        # 设置输出文件路径
        output_path = "yuzhong_roads.geojson"

        # 转换并保存为GeoJSON
        print("正在转换为GeoJSON格式...")
        graph_to_geojson(G, output_path)

        # 输出基本统计信息
        print("\n道路网络统计信息:")
        print(f"节点数量: {len(G.nodes)}")
        print(f"道路段数量: {len(G.edges)}")
    else:
        print("获取道路网络失败")


if __name__ == "__main__":
    main()